package org.example.ratelimit;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 计数器限流
 * 单位时间内，请求处理的最大数量是固定的
 */
@Slf4j
public class CounterLimiter {

    // 起始时间
    private static long startTime = System.currentTimeMillis();

    // 1秒之内最大请求数量为2
    private static long interval = 1000;
    private static long maxCount = 2;

    private static AtomicLong counter = new AtomicLong(0);

    private static long tryAccquire(long taskId, int turn) {
        long now = System.currentTimeMillis();
        // 区间之内
        if ((now - startTime) < interval) {
            long count = counter.incrementAndGet();
            if (count <= maxCount) {
                return count;
            } else {
                return -count;
            }
        } else {
            // 没有拿到令牌
            synchronized (CounterLimiter.class) {
                log.info("新时间区到了，taskId：{}, turns：{}", taskId, turn);
                // 再一次判断，防止重复初始化
                if ((now - startTime) > interval) {
                    counter.set(0);
                    startTime = now;
                }
            }
            return 0;
        }
    }

    public static void main(String[] args) throws Exception{
        ExecutorService service = Executors.newFixedThreadPool(2);
        CountDownLatch latch = new CountDownLatch(2);

        // 每个线程执行20次
        final int turns = 20;
        final AtomicInteger failed = new AtomicInteger();
        long start = System.currentTimeMillis();
        for (int i = 0; i < 2; i++) {
            service.submit(() -> {
                try {
                    for (int j = 0; j < turns; j++) {
                        long taskId = Thread.currentThread().getId();
                        long count = tryAccquire(taskId, j);
                        if (count <= 0) {
                            failed.incrementAndGet();
                        }
                        Thread.sleep(200);
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    latch.countDown();
                }
            });
        }

        latch.await();
        log.info("总的请求数：{}", turns * 2);
        log.info("被限制的请求数：{}", failed.get());
        log.info("成功的请求数：{}", turns * 2 - failed.get());
        log.info("成功比例：{}", (float)((turns * 2 - failed.get())) / (float) ((turns * 2)));
        log.info("请求总耗时：{}秒", (System.currentTimeMillis() - start) / 1000F);
    }
}
